# 邓迪心数据集v1
from typing import Dict, Tuple, List
import os
import shutil
import random
import numpy as np
import matplotlib.pyplot as plt
import torch
from core.lrp_config import LrpConfig as LC

class DxdDataV1(object):
    @staticmethod
    def prepare_ds(params:Dict = {}) -> None:
        '''
        准备数据集
        '''
        rec_db, dev_recs, raw_datas = DxdDataV1.load_raw_datas(params={})
        did = DxdDataV1.get_max_recs_dev(dev_recs=dev_recs)
        recs = DxdDataV1.get_rid_recs(did, dev_recs, raw_datas)
        norm_recs = DxdDataV1.generate_norm_recs_dict(recs)
        for k, v in norm_recs.items():
            recs[k] = v
        frames = DxdDataV1.extract_frames(recs[rid]['data'], recs[rid]['label'])
        seq = 0
        DxdDataV1.save_frames(seq, did, recs)

    @staticmethod
    def load_raw_datas(params:Dict = {}) -> Tuple[Dict, Dict, Dict]:
        # 1.1. 读入：{rec_id, {数据记录描述信息：开始时间、地点......}}
        csv_fn = 'work/datas/v0/time_periods.csv'
        rec_db = DxdDataV1.get_rec_db(csv_fn=csv_fn)
        # 1.2. 读入: {'devid': {'tls': [雷电原始数据ID列表], 'norms': [正常数据ID列表]}}
        dev_recs = DxdDataV1.get_devs_recs(csv_fn=csv_fn)
        # 2.1. 读入原始数据: {rec_id, ndarray}
        electric_fn = 'work/datas/v0/electric.txt'
        raw_datas = DxdDataV1.get_raw_datas(electric_fn=electric_fn)
        return rec_db, dev_recs, raw_datas
    
    @staticmethod
    def get_rec_db(csv_fn:str) -> Dict:
        '''
        获取记录数据基本信息
        '''
        rec_db = {}
        first_row = True
        with open(csv_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                if first_row:
                    first_row = False
                    continue
                arrs1 = row.strip().split(',')
                if len(arrs1) <= 1:
                    continue
                rec_id = arrs1[LC.IDV1_id]
                rec_db[rec_id] = {
                    'id': arrs1[LC.IDV1_id],
                    'start_time': arrs1[LC.IDV1_start_time],
                    'did': arrs1[LC.IDV1_did],
                    'dname': arrs1[LC.IDV1_dname],
                    'device_lng': arrs1[LC.IDV1_device_lng],
                    'device_lat': arrs1[LC.IDV1_device_lat],
                    'max_elevel': arrs1[LC.IDV1_max_elevel],
                    'warning_start_time': arrs1[LC.IDV1_warning_start_time],
                    'warning_end_time': arrs1[LC.IDV1_warning_end_time],
                    'first_lightning_time': arrs1[LC.IDV1_warning_end_time],
                    'first_lightning_lng': arrs1[LC.IDV1_first_lightning_lng],
                    'first_lightning_lat': arrs1[LC.IDV1_first_lightning_lat],
                    'first_lightning_distance': arrs1[LC.IDV1_first_lightning_distance],
                    'first_lightning_strength': arrs1[LC.IDV1_first_lightning_distance],
                    'first_lw_level': arrs1[LC.IDV1_first_lw_level],
                    'last_lightning_time': arrs1[LC.IDV1_last_lightning_time],
                    'last_lightning_lng': arrs1[LC.IDV1_last_lightning_lng],
                    'last_lightning_lat': arrs1[LC.IDV1_last_lightning_lat],
                    'last_lightning_distance': arrs1[LC.IDV1_last_lightning_distance],
                    'last_lightning_strength': arrs1[LC.IDV1_last_lightning_strength],
                    'last_lw_level': arrs1[LC.IDV1_last_lw_level]
                }
        return rec_db
    
    @staticmethod
    def get_devs_recs(csv_fn:str) -> Dict:
        '''
        求出所有设备雷电记录条数和正常记录条数
        返回值：Dict
        {
            'devid': {
                'tls': [1, 2, 3, ...],
                'norms': [11, 22, 33, ...]
            }
        }
        '''
        # csv_fn = 'work/tlp/datas/v1/time_periods.csv'
        devs = {}
        max_num = 0
        with open(csv_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                arrs0 = row.strip().split(',')
                if len(arrs0) <= 1:
                    continue
                rec_id = arrs0[LC.IDV1_id]
                did = arrs0[LC.IDV1_did]
                if did not in devs:
                    if arrs0[LC.IDV1_start_time] == '':
                        devs[did] = {
                            'tls': [],
                            'norms': [rec_id]
                        }
                    else:
                        devs[did] = {
                            'tls': [rec_id],
                            'norms': []
                        }
                else:
                    if arrs0[LC.IDV1_start_time] == '':
                        devs[did]['norms'].append(rec_id)
                    else:
                        devs[did]['tls'].append(rec_id)
        return devs
    
    @staticmethod
    def get_raw_datas(electric_fn:str) -> Dict:
        '''
        从electric.txt文件中获取所有记录，格式：
        {
            rec_id: [1.0, 0.0, 2.0],
            ......
        }
        '''
        raw_datas = {}
        num = 0

        with open(electric_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                arrs0 = row.strip().split(' ')
                rec_id = arrs0[0]
                arrs1 = arrs0[1].split(',')
                rec = torch.tensor([float(xi) for xi in arrs1])
                raw_datas[rec_id] = rec
                num += 1
                if num % 1000 == 0:
                    print(f'处理完成{num}条数据......')
        return raw_datas
    
    @staticmethod
    def get_max_recs_dev(dev_recs:Dict) -> str:
        # 按照数据记录条数从多到少排充 did=F56AC57A1181CD
        dids = sorted(dev_recs.keys(), key=lambda dev_id: len(dev_recs[dev_id]["tls"]) + len(dev_recs[dev_id]["norms"]), reverse=True)
        return dids[0]
    
    @staticmethod
    def get_recs(did:str, dev_recs:Dict, raw_datas:Dict) -> Tuple[List, List]:
        '''
        所有原始记录的列表：
        recs: [ndarray, ndarray, ....]
        labels: [0, 1, ......] 0-代表正常; 1-代表雷电;
        '''
        recs, labels = [], []
        tls_ids = dev_recs[did]['tls']
        for rid in tls_ids:
            data = raw_datas[rid]
            recs.append(data)
            labels.append(1)
        norm_ids = dev_recs[did]['norms']
        for rid in norm_ids:
            data = raw_datas[rid]
            recs.append(data)
            labels.append(0)
        return recs, labels
    
    @staticmethod
    def get_rid_recs(did:str, dev_recs:Dict, raw_datas:Dict) -> Dict:
        '''
        所有原始记录的列表：
        {
            'rid': {
                'data': [1.0, 2.0, ..., 1.0],
                'label': 0 # 0-代表正常; 1-代表雷电;
            }
        }
        '''
        recs = {}
        tls_ids = dev_recs[did]['tls']
        for rid in tls_ids:
            recs[rid] = {
                'data': raw_datas[rid],
                'label': 1
            }
        norm_ids = dev_recs[did]['norms']
        for rid in norm_ids:
            recs[rid] = {
                'data': raw_datas[rid],
                'label': 0
            }
        return recs
    
    @staticmethod
    def extract_frames(data:torch.Tensor, data_type:int = 1) -> torch.Tensor:
        '''
        从原始数据记录中抽取数据帧，对于雷电记录，取到第一次雷电开始后10分钟，对于正常记录取全部长度
        '''
        seq = 0
        X = None
        if data_type == 1:
            cnt = 90 * 60 + 10*60 # 对于雷电记录，取到第一次雷电后10分钟
        else:
            cnt = len(data) # 对于正常记录取全部数据
        for ptr in range(0, cnt - LC.C_frame_size + 1, LC.C_stride):
            seq += 1
            frame = data[ptr : ptr+LC.C_frame_size]
            if X is None:
                X = frame
            else:
                X = torch.vstack((X, frame))
        return X
    
    @staticmethod
    def generate_norm_recs_dict(recs:Dict) -> Dict:
        '''
        由雷电记录生成正常记录，将后90分钟数据加上前10分钟数据组成正常帧
        '''
        norm_recs = {}
        for k, v in recs.items():
            data = torch.hstack((v['data'][-90*60:], v['data'][:10*60]))
            print(f'### data: {data.shape};')
            norm_recs[f'n{k}'] = {
                'data': data,
                'label': 0
            }
        return norm_recs
    
    @staticmethod
    def save_frames(seq:int, did:str, recs:Dict) -> int:
        '''
        recs[rid]['data']保存为d_**.pt文件，在dss.txt文件中添加一行：
            数据文件名,设备编号,记录编号,记录类别,位置,类别
            work/datasets/dhlp/v1/d00/d00/d00/d00/d00/d_00.pt,did,rid,recs[rid]['label'],1,0
        其中：
            最后一位：
            倒数第二位：根据idx，其他一律为0
                # 50~59：距离雷电发生1~10分钟的帧
                # 40~49：距离雷电发生11~20分钟
                # 30~39：距离雷电发生21~30分钟
        '''
        cls = 0
        ds_txt_fn = './work/datasets/dhlp/v1/ds.txt'
        with open(ds_txt_fn, 'w', encoding='utf-8') as wfd:
            for rid, item in recs.items():
                label = item['label']
                data = item['data']
                frames = DxdDataV1.extract_frames(data, label)
                cnt = frames.shape[0]
                for idx in range(cnt):
                    if idx>=50 and idx<=59:
                        pos = 1
                    elif idx>=40 and idx<=49:
                        pos = 2
                    elif idx>=30 and idx<=39:
                        pos = 3
                    else:
                        pos = 0
                    fn = DxdDataV1.generate_fn(seq)
                    seq += 1
                    if seq % 100 == 0:
                        print(f'    保存{seq}帧......')
                    torch.save(frames[idx], fn)
                    wfd.write(f'{fn},{did},{rid},{label},{pos},{idx},{cls}\n')
        return seq

    @staticmethod
    def generate_fn(seq:int) -> str:
        '''
        根据seq数值，生成如下格式文件名：work/datasets/dhlp/v1/d00/d00/d00/d00/d00/d_00.pt，并创建数据集目录
        '''
        fn = ''
        raw_str = f'{seq:012d}'
        folder = 'work/datasets/dhlp/v1'
        for i in range(6):
            folder = f'{folder}/d{raw_str[i*2:(i+1)*2]}'
            os.makedirs(folder, exist_ok=True) # 目录存在时不会报异常
        fn = f'{folder}.pt'
        return fn
    
    
    
    @staticmethod
    def generate_norm_recs(recs:List) -> Tuple[List, List]:
        norm_recs, norm_labels = [], []
        cnt = len(recs)
        for rec in recs:
            left = rec[0:60*60]
            print(f'### {type(rec)}; left: {type(left)};')
            rv = random.randint(0, cnt-1)
            right = recs[rv][-60*60:]
            new_rec = torch.hstack((left, right))
            norm_recs.append(new_rec)
            norm_labels.append(0)
        return norm_recs, norm_labels
    
    @staticmethod
    def split_train_test(recs:List, labels:List) -> Tuple[List, List, List, List]:
        '''
        将数据集划分为训练集和测试集
        '''
        seed = 102
        rvo = random.Random(seed)
        train_datas, test_datas = [], []
        train_labels, test_labels = [], []
        cnt = len(recs)
        for idx in range(cnt):
            rv = rvo.random()
            if rv > 0.9:
                test_datas.append(recs[idx])
                test_labels.append(labels[idx])
            else:
                train_datas.append(recs[idx])
                train_labels.append(labels[idx])
        return train_datas, train_labels, test_datas, test_labels
    





    # ds.txt每行中字段索引值
    DSF_FN, DSF_DID, DSF_RID, DSF_LABEL, DSF_POS, DSF_IDX, DSF_CLS = 0, 1, 2, 3, 4, 5, 6
    @staticmethod
    def load_ds_txt(ds_txt_fn:str) -> Tuple[List, Dict, Dict, Dict, Dict]:
        '''
        解析ds.txt文件，返回：
            List：rid列表
            Dict：[rid, frame]
        '''
        rid_set = set()
        # 分别代表正常帧、雷电前10分钟、雷电前10~20分钟、雷电前20~30分钟记录编号与帧的字典，结构为：
        # {rid: [{'fn': '*.pt', 'label': 0/1, 'cls':0}, ...]}
        rid_frame0s, rid_frame1s, rid_frame2s, rid_frame3s = {}, {}, {}, {}
        with open(ds_txt_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                rid_set.add(arrs0[DxdDataV1.DSF_RID])
                item = {
                    'fn': arrs0[DxdDataV1.DSF_FN],
                    'label': int(arrs0[DxdDataV1.DSF_LABEL]),
                    'cls': int(arrs0[DxdDataV1.DSF_CLS])
                }
                if arrs0[DxdDataV1.DSF_POS] == '0':
                    if arrs0[DxdDataV1.DSF_RID] not in rid_frame0s:
                        rid_frame0s[arrs0[DxdDataV1.DSF_RID]] = [item]
                    else:
                        rid_frame0s[arrs0[DxdDataV1.DSF_RID]].append(item)
                elif arrs0[DxdDataV1.DSF_POS] == '1':
                    if arrs0[DxdDataV1.DSF_RID] not in rid_frame1s:
                        rid_frame1s[arrs0[DxdDataV1.DSF_RID]] = [item]
                    else:
                        rid_frame1s[arrs0[DxdDataV1.DSF_RID]].append(item)
                elif arrs0[DxdDataV1.DSF_POS] == '2':
                    if arrs0[DxdDataV1.DSF_RID] not in rid_frame2s:
                        rid_frame2s[arrs0[DxdDataV1.DSF_RID]] = [item]
                    else:
                        rid_frame2s[arrs0[DxdDataV1.DSF_RID]].append(item)
                elif arrs0[DxdDataV1.DSF_POS] == '3':
                    if arrs0[DxdDataV1.DSF_RID] not in rid_frame3s:
                        rid_frame3s[arrs0[DxdDataV1.DSF_RID]] = [item]
                    else:
                        rid_frame3s[arrs0[DxdDataV1.DSF_RID]].append(item)
        return list(rid_set), rid_frame0s, rid_frame1s, rid_frame2s, rid_frame3s
    
    @staticmethod
    def show_red_alarm_frames(ds_txt_fn:str, did:str, rid:str, mode:int = 1) -> None:
        '''
        可视化指定设备和记录雷电前10分钟数据帧，所括：三种类型的10张png文件，一个显示10个图片的静态网页，三种类型为：原始数据、FFT变换数据、希尔伯特黄变换
        mode: 1---原始数据；2---FFT后数据；
        '''
        frames_raw = []
        with open(ds_txt_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                ds_did = arrs0[DxdDataV1.DSF_DID]
                ds_rid = arrs0[DxdDataV1.DSF_RID]
                item = {
                    'fn': arrs0[DxdDataV1.DSF_FN],
                    'label': int(arrs0[DxdDataV1.DSF_LABEL]),
                    'cls': int(arrs0[DxdDataV1.DSF_CLS])
                }
                if ds_did==did and ds_rid==rid and arrs0[DxdDataV1.DSF_POS]=='1': # 雷电前10分钟的数据帧
                    pt_fn = arrs0[0]
                    data = torch.load(pt_fn, weights_only=True)
                    frames_raw.append({
                        'fn': pt_fn,
                        'data': data,
                        'cls': arrs0[DxdDataV1.DSF_CLS]
                    })
        folder = f'./work/datas/dhlp/{did}'
        if not os.path.exists(folder):
            os.makedirs(folder, exist_ok=True)
        folder = f'{folder}/{rid}'
        if not os.path.exists(folder):
            os.makedirs(folder, exist_ok=True)
        raw_folder = f'{folder}/raw'
        if not os.path.exists(raw_folder):
            os.makedirs(raw_folder, exist_ok=True)
        fft_folder = f'{folder}/fft'
        if not os.path.exists(fft_folder):
            os.makedirs(fft_folder, exist_ok=True)
        cnt = len(frames_raw)
        if mode == 1:
            with open(f'{raw_folder}/index.txt', 'w', encoding='utf-8') as wfd:
                for idx in range(cnt):
                    y_np = frames_raw[idx]['data'].numpy()
                    x = np.arange(y_np.shape[0])
                    plt.plot(x, y_np)
                    plt.title(f'type: {frames_raw[idx]["cls"]}')
                    plt.savefig(f'{raw_folder}/frm_{idx:03d}.png')
                    plt.close()
                    wfd.write(f'{frames_raw[idx]["fn"]},frm_{idx:03d}.png\n')
        else:
            with open(f'{fft_folder}/index.txt', 'w', encoding='utf-8') as wfd:
                for idx in range(cnt):
                    y_np = frames_raw[idx]['data'].numpy()
                    Y = np.fft.fft(y_np)
                    N = y_np.shape[0]
                    frequencies = np.fft.fftfreq(N)
                    # x = np.arange(y_np.shape[0])
                    plt.plot(frequencies, np.abs(Y))
                    plt.title(f'type: {frames_raw[idx]["cls"]}')
                    plt.savefig(f'{fft_folder}/frm_{idx:03d}.png')
                    plt.close()
                    wfd.write(f'{frames_raw[idx]["fn"]},frm_{idx:03d}.png\n')

    @staticmethod
    def show_all_red_alarm_frames(ds_txt_fn:str, mode:int = 1) -> None:
        '''
        显示所有记录的雷电前10分钟帧
        '''
        num = 0
        with open(ds_txt_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                did = arrs0[DxdDataV1.DSF_DID]
                rid = arrs0[DxdDataV1.DSF_RID]
                if arrs0[DxdDataV1.DSF_POS]=='1' and arrs0[DxdDataV1.DSF_LABEL]=='1': # 雷电前10分钟的数据帧
                    DxdDataV1.show_red_alarm_frames(ds_txt_fn, did, rid, mode=mode)
                num += 1
                if num % 1000 == 0:
                    print(f'处理完成{num}帧，进度：{num/53960.0:02f};')
